import torch  # 导入 PyTorch 库
import torch.nn as nn  # 导入 PyTorch 的神经网络模块
import torchvision  # 导入 torchvision 库，包含图像处理相关的工具和数据集
import torchvision.transforms as transforms  # 导入 torchvision 的 transforms 模块，用于图像预处理

# 设备配置：如果 CUDA 可用，则使用 GPU；否则使用 CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 超参数设置
sequence_length = 28  # 输入序列的长度（MNIST 图像的高度）
input_size = 28  # 输入特征的维度（MNIST 图像的宽度）
hidden_size = 128  # RNN 隐藏层的大小
num_layers = 2  # RNN 的层数
num_classes = 10  # 输出类别的数量（0-9 十个数字）
batch_size = 100  # 批次大小
num_epochs = 4  # 训练的总轮数 (epochs)
learning_rate = 0.005  # 学习率

# 加载 MNIST 数据集
train_dataset = torchvision.datasets.MNIST(root='./data/',  # 数据集存放的路径
                                           train=True,  # 加载训练集
                                           transform=transforms.ToTensor(),  # 将图像转换为 Tensor，并归一化到 [0, 1]
                                           download=True)  # 如果本地没有数据，则下载

test_dataset = torchvision.datasets.MNIST(root='./data/',  # 数据集存放的路径
                                          train=False,  # 加载测试集
                                          transform=transforms.ToTensor())  # 将图像转换为 Tensor，并归一化到 [0, 1]

# 创建数据加载器 (DataLoader)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,  # 训练数据集
                                           batch_size=batch_size,  # 批次大小
                                           shuffle=True)  # 打乱训练数据

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,  # 测试数据集
                                          batch_size=batch_size,  # 批次大小
                                          shuffle=False)  # 测试数据不需要打乱


# 定义传统的 RNN 模型
class RNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(RNN, self).__init__()  # 调用父类的初始化方法
        self.hidden_size = hidden_size  # 隐藏层大小
        self.num_layers = num_layers  # RNN 层数
        # 创建 RNN 层 (nn.RNN)
        self.rnn = nn.RNN(input_size, hidden_size, num_layers, batch_first=True, nonlinearity='relu')
        # nonlinearity='relu' 指定激活函数为 ReLU
        self.fc = nn.Linear(hidden_size, num_classes)  # 全连接层，将 RNN 输出映射到类别

    def forward(self, x):
        # 初始化隐藏状态 h0 (全零)
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)

        # 前向传播 RNN
        out, _ = self.rnn(x, h0)  # out: (batch_size, sequence_length, hidden_size)

        # 只取最后一个时间步的输出，并送入全连接层
        out = self.fc(out[:, -1, :])
        return out


# 实例化 RNN 模型，并将其移动到指定设备 (CPU 或 GPU)
model = RNN(input_size, hidden_size, num_layers, num_classes).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数，用于多分类问题
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  # Adam 优化器

# 训练模型
total_step = len(train_loader)  # 总的训练步数（批次数）
for epoch in range(num_epochs):  # 遍历每个 epoch
    for i, (images, labels) in enumerate(train_loader):  # 遍历每个批次的数据
        # 将图像数据和标签移动到指定设备
        images = images.reshape(-1, sequence_length, input_size).to(device)
        labels = labels.to(device)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)  # 计算损失

        # 反向传播和优化
        optimizer.zero_grad()  # 梯度清零
        loss.backward()  # 反向传播，计算梯度

        # 梯度裁剪，防止梯度爆炸 (最大范数设置为 5)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=5)

        optimizer.step()  # 更新模型参数

        # 每 100 个批次打印一次训练信息
        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                  .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))

# 测试模型
with torch.no_grad():  # 在测试阶段，不计算梯度
    correct = 0
    total = 0
    for images, labels in test_loader:  # 遍历测试集的每个批次
        images = images.reshape(-1, sequence_length, input_size).to(device)
        labels = labels.to(device)
        outputs = model(images)  # 前向传播
        _, predicted = torch.max(outputs.data, 1)  # 获取预测的类别
        total += labels.size(0)  # 更新总样本数
        correct += (predicted == labels).sum().item()  # 更新预测正确的样本数

    # 打印模型在测试集上的准确率
    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))
